const std = @import("std");
const dtb = @import("dtb");
const os = @import("os.zig");
pub const Node = dtb.Node;

var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator);
var fdt: *Node = undefined;
var fdt_addr: usize = undefined;

const Nodes = std.ArrayList(*const Node);
var device_map: std.StringArrayHashMap(Nodes) = .init(arena.allocator());

/// Initializes the DTB parser using the FDT address in memory.
pub fn init(addr: usize) !void {
    fdt_addr = addr;
    const raw_size: *u32 = @ptrFromInt(fdt_addr + 4);
    const fdt_size: u32 = std.mem.bigToNative(u32, raw_size.*);
    const fdt_bytes: [*]u8 = @ptrFromInt(fdt_addr);

    const allocator = arena.allocator();
    fdt = try dtb.parse(allocator, fdt_bytes[0..fdt_size]);
}

pub fn deinit() void {
    // Reclaim FDT memory region after driver initialization is complete
    // The dtb memory is no longer needed and can be returned to allocator
    const fdt_addr_aligned = std.mem.alignBackward(usize, fdt_addr, os.PGSIZE);
    const mem_base, const mem_size = memory();
    const heap_end = std.mem.alignBackward(usize, @intCast(mem_base + mem_size), os.PGSIZE);
    os.mem.PageAllocator.grow(fdt_addr_aligned, heap_end);

    // Cleans up all allocated resources used by the DTB parser and device mapping.
    arena.deinit();
}

/// Returns the number of CPU nodes found in the device tree.
pub fn ncpus() usize {
    var n: usize = 0;
    const cpus = fdt.child("cpus").?;
    for (cpus.children) |c| {
        if (match(c, "cpu")) n += 1;
    }
    return n;
}

/// Retrieves the memory regions defined in the device tree.
pub fn memory() [2]u128 {
    const mem_node = child(fdt, "memory").?;
    return mem_node.prop(.Reg).?[0];
}

/// Gets the stdout UART node from the chosen field.
pub fn stdout() ?*const Node {
    const chosen = fdt.child("chosen").?;
    const stdout_path = for (chosen.props) |p| {
        if (p == .Unknown and std.mem.eql(u8, p.Unknown.name, "stdout-path")) {
            break p.Unknown.value[0 .. p.Unknown.value.len - 1];
        }
    } else unreachable;

    return find(fdt, stdout_path);
}

/// Finds the interrupt controller node in the SOC subtree.
pub fn intc() ?*const Node {
    const soc = fdt.child("soc").?;
    for (soc.children) |ch| {
        ch.prop(.InterruptController) orelse continue;
        return ch;
    }
    return null;
}

/// Extracts the base address from the node's register property.
pub fn base(node: *const Node) usize {
    return @intCast(node.prop(.Reg).?[0][0]);
}

/// Parses the FDT and maps device memory regions into kernel virtual address space.
pub fn probe() void {
    errdefer |err| @panic(@errorName(err));

    const soc = fdt.child("soc").?;
    for (soc.children) |ch| {
        // Skip nodes without compatible property
        const compatibles = ch.prop(.Compatible) orelse continue;
        const dev = os.driver.find(compatibles) orelse continue;
        const reg = ch.prop(.Reg) orelse continue;

        // Map device memory into kernel virtual address space
        const address: usize = @intCast(reg[0][0]);
        const size: usize = @intCast(reg[0][1]);
        try os.mem.kvmmap(address, size, .{ .readable = true, .writable = true });

        // Add device node to global device map, creating entry if new
        const gop = try device_map.getOrPut(std.mem.span(dev.name));
        if (!gop.found_existing) {
            gop.value_ptr.* = .empty;
        }
        try gop.value_ptr.append(device_map.allocator, ch);
    }
}

/// Returns device nodes matching the given compatible string, or null if none found.
pub fn get(compatible: []const u8) ?Nodes {
    return device_map.get(compatible);
}

/// Finds a node by traversing the full path from a starting node.
fn find(start: *const Node, path: []const u8) ?*const Node {
    const PathIterator = std.fs.path.ComponentIterator(.posix, u8);
    var it: PathIterator = .init(path);

    var node: *const Node = start;
    while (it.next()) |component| {
        node = child(node, component.name) orelse return null;
    }
    return node;
}

/// Returns a direct child node of the given node by name.
fn child(node: *const Node, child_name: []const u8) ?*Node {
    for (node.children) |c| {
        if (match(c, child_name)) return c;
    }
    return null;
}

/// Checks if a node's name matches the given name, allowing omitted unit addresses.
fn match(node: *const Node, name: []const u8) bool {
    if (node.name.len < name.len)
        return false;

    if (!std.mem.eql(u8, node.name[0..name.len], name))
        return false;

    if (node.name.len == name.len)
        return true;

    if (node.name[name.len] == 0)
        return true;

    if (node.name[name.len] == '@')
        return true;

    return false;
}
